from math import sqrt

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression


# 读取 CSV 文件
file_path = 'Landset8_20240925.csv'  # 替换为你的 CSV 文件路径
df = pd.read_csv(file_path)

# 提取时间
date = file_path.split('_')[1].split('.')[0]  # 假设时间在第二个下划线分隔的部分
print(f"提取的时间为: {date}")

# 提取卫星名（假设卫星名在第一个下划线分隔的部分）
satellite_name = file_path.split('_')[0]
print(f"卫星名为: {satellite_name}")

def get_wet_coefficients(satellite_name):
    if satellite_name == "Landset6":
        return {"Blue": 0.0315, "Green": 0.2021, "Red": 0.3012, "NIR": 0.1594, "SWIR1": -0.6806, "SWIR2": -0.6109}
    elif satellite_name == "Landset7":
        return {"Blue": 0.2626, "Green": 0.2141, "Red": 0.0926, "NIR": 0.0656, "SWIR1": -0.7629, "SWIR2": -0.5388}
    else:
        return {"Blue": 0.1511, "Green": 0.1973, "Red": 0.3283, "NIR": 0.3407, "SWIR1": -0.7117, "SWIR2": -0.4559}

for col in df.columns:
    if col in ['SR_B1', 'SR_B2', 'SR_B3', 'SR_B4', 'SR_B5', 'SR_B6', 'SR_B7']:
        df[col] = df[col] * 0.0000275 - 0.2
    elif col == 'ST_B10' or col == 'ST_B6':
        df[col] = df[col] * 0.00341802 + 149.0

# 定义新列名
new_column_names = {
    'SR_B2': 'Blue',
    'SR_B3': 'Green',
    'SR_B4': 'Red',
    'SR_B5': 'NIR',
    'SR_B6': 'SWIR1',
    'SR_B7': 'SWIR2',
    'ST_B10': 'TIR'
}

# 重命名列
df.rename(columns=new_column_names, inplace=True)

wet_coefficients = get_wet_coefficients(satellite_name)

# 计算生态指数
df['WET'] = (df['Blue'] * wet_coefficients['Blue'] +
                  df['Green'] * wet_coefficients['Green'] +
                  df['Red'] * wet_coefficients['Red'] +
                  df['NIR'] * wet_coefficients['NIR'] +
                  df['SWIR1'] * wet_coefficients['SWIR1'] +
                  df['SWIR2'] * wet_coefficients['SWIR2'])

df['SI'] = (df['SWIR1'] + df['Red'] - df['NIR'] - df['Blue']) / \
                (df['SWIR1'] + df['Red'] + df['NIR'] + df['Blue'])

df['IBI'] = ((2.0 * df['SWIR1']) / (df['SWIR1'] + df['NIR']) -
                  (df['NIR'] / (df['NIR'] + df['Red']) +
                   df['Green'] / (df['Green'] + df['SWIR1']))) / \
                 ((2.0 * df['SWIR1']) / (df['SWIR1'] + df['NIR']) +
                  (df['NIR'] / (df['NIR'] + df['Red']) +
                   df['Green'] / (df['Green'] + df['SWIR1'])))

df['NDBSI'] = (df['IBI'] + df['SI']) / 2
df['NDVI'] = (df['NIR'] -df['Red']) / (df['NIR'] + df['Red'])
df['LST'] = df['TIR'] - 273.15  # 假设TIR列已经重命名为'TIR'
df['Albedo'] = 0.356 * df['Blue'] + 0.130 * df['Red'] + 0.373 * df['NIR'] + 0.072 * df['SWIR2'] - 0.1108
# 使用线性回归计算 NDVI 和 Albedo 之间的斜率
# 将 NDVI 数据转换为二维数组
NDVI = df['NDVI'].values.reshape(-1, 1)
Albedo = df['Albedo'].values
# 创建线性回归模型
model = LinearRegression()
# 拟合模型
model.fit(NDVI, Albedo)
# 获取斜率 a
K = model.coef_[0]
df['DDI'] = -1 * (1 / K) * df['NDVI'] - df['Albedo']
df['DI'] = -1 * df['DDI']

df['PMDI'] = df['Red'] - df['NIR']

df['SI3'] = np.sqrt(df['Green'] * df['Green'] + df['Red'] * df['Red'])
df['NDSI'] = (df['Red'] - df['NIR']) / (df['Red'] + df['NIR'])
df['SI_T'] = (df['Red'] / df['NIR']) * 100
df['CSI'] = (df['SI_T'] + df['NDSI'] + df['SI3']) / 3


# 保存需要的列到新的文件
output_columns = ['longitude', 'latitude', 'NDVI', 'WET', 'NDBSI', 'LST', 'PMDI', 'DI', 'CSI']
output_file_path = 'output.csv'  # 输出文件路径
df[output_columns].to_csv(output_file_path, index=False)
